package com.amu.aidemo.android.zhipu

import android.util.Log
import com.google.gson.Gson
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import okhttp3.Response
import okhttp3.sse.EventSource
import okhttp3.sse.EventSourceListener
import okhttp3.sse.EventSources
import java.util.concurrent.TimeUnit

private const val TAG = "ZhipuAiClient"

/**
 * Zhipu AI Client for Chat Completion API
 * 
 * API Documentation: https://open.bigmodel.cn/dev/api
 * Endpoint: https://open.bigmodel.cn/api/paas/v4/chat/completions
 */
class ZhipuAiClient(
    private val apiKey: String
) {
    private val baseUrl = "https://open.bigmodel.cn/api/paas/v4"
    private val gson = Gson()
    
    private val client = OkHttpClient.Builder()
        .connectTimeout(30, TimeUnit.SECONDS)
        .readTimeout(60, TimeUnit.SECONDS)
        .writeTimeout(60, TimeUnit.SECONDS)
        .build()
    
    /**
     * Callback for streaming response
     */
    var onStreamChunk: ((String) -> Unit)? = null
    var onStreamComplete: (() -> Unit)? = null
    var onError: ((String) -> Unit)? = null
    
    // Guard flag to prevent duplicate completion calls
    private var isStreamCompleted = false
    
    /**
     * Send chat completion request with streaming
     */
    suspend fun chatCompletionStream(
        messages: List<ZhipuMessage>,
        model: String = "glm-4.6",
        temperature: Float = 0.8f,
        maxTokens: Int = 2048
    ) = withContext(Dispatchers.IO) {
        try {
            val startTime = System.currentTimeMillis()
            Log.d(TAG, "⏱️ [START] Chat completion request started")
            
            // Reset completion flag for new stream
            isStreamCompleted = false
            
            val request = ZhipuChatRequest(
                model = model,
                messages = messages,
                stream = true,
                temperature = temperature,
                maxTokens = maxTokens
            )
            
            val requestBody = gson.toJson(request)
                .toRequestBody("application/json".toMediaType())
            
            val httpRequest = Request.Builder()
                .url("$baseUrl/chat/completions")
                .addHeader("Authorization", "Bearer $apiKey")
                .addHeader("Content-Type", "application/json")
                .post(requestBody)
                .build()
            
            val requestPreparedTime = System.currentTimeMillis()
            Log.d(TAG, "⏱️ [PREPARED] Request prepared in ${requestPreparedTime - startTime}ms")
            Log.d(TAG, "Model: $model, Messages: ${messages.size}")
            
            var firstChunkTime: Long? = null
            var chunkCount = 0
            
            val eventSourceListener = object : EventSourceListener() {
                override fun onOpen(eventSource: EventSource, response: Response) {
                    val connectionTime = System.currentTimeMillis()
                    Log.d(TAG, "⏱️ [CONNECTED] SSE connection opened in ${connectionTime - requestPreparedTime}ms")
                    Log.d(TAG, "⏱️ [TOTAL TO CONNECT] ${connectionTime - startTime}ms from start")
                }
                
                override fun onEvent(
                    eventSource: EventSource,
                    id: String?,
                    type: String?,
                    data: String
                ) {
                    if (data == "[DONE]") {
                        if (!isStreamCompleted) {
                            isStreamCompleted = true
                            Log.d(TAG, "Stream completed")
                            onStreamComplete?.invoke()
                        }
                        return
                    }
                    
                    try {
                        val chunk = gson.fromJson(data, ZhipuChatResponse::class.java)
                        val content = chunk.choices.firstOrNull()?.delta?.content
                        
                        if (!content.isNullOrEmpty()) {
                            chunkCount++
                            if (firstChunkTime == null) {
                                firstChunkTime = System.currentTimeMillis()
                                Log.d(TAG, "⏱️ [FIRST CHUNK] First content received in ${firstChunkTime!! - startTime}ms")
                            }
                            onStreamChunk?.invoke(content)
                        }
                        
                        // Check if finished
                        val finishReason = chunk.choices.firstOrNull()?.finishReason
                        if ((finishReason == "stop" || finishReason == "length") && !isStreamCompleted) {
                            isStreamCompleted = true
                            val endTime = System.currentTimeMillis()
                            Log.d(TAG, "⏱️ [FINISHED] Stream finished: $finishReason")
                            Log.d(TAG, "⏱️ [STATS] Total time: ${endTime - startTime}ms, Chunks: $chunkCount")
                            onStreamComplete?.invoke()
                        }
                    } catch (e: Exception) {
                        Log.e(TAG, "Failed to parse SSE data: $data", e)
                    }
                }
                
                override fun onClosed(eventSource: EventSource) {
                    Log.d(TAG, "SSE connection closed")
                    if (!isStreamCompleted) {
                        isStreamCompleted = true
                        onStreamComplete?.invoke()
                    }
                }
                
                override fun onFailure(
                    eventSource: EventSource,
                    t: Throwable?,
                    response: Response?
                ) {
                    isStreamCompleted = true  // Mark as completed on failure
                    val errorMsg = t?.message ?: "Unknown error"
                    Log.e(TAG, "SSE connection failed: $errorMsg", t)
                    
                    // Try to parse error response
                    response?.body?.string()?.let { body ->
                        try {
                            val errorResponse = gson.fromJson(body, ZhipuErrorResponse::class.java)
                            onError?.invoke("${errorResponse.error.code}: ${errorResponse.error.message}")
                        } catch (e: Exception) {
                            onError?.invoke(errorMsg)
                        }
                    } ?: onError?.invoke(errorMsg)
                }
            }
            
            EventSources.createFactory(client)
                .newEventSource(httpRequest, eventSourceListener)
            
        } catch (e: Exception) {
            Log.e(TAG, "Failed to send request", e)
            onError?.invoke(e.message ?: "Unknown error")
        }
    }
    
    /**
     * Send chat completion request without streaming
     */
    suspend fun chatCompletion(
        messages: List<ZhipuMessage>,
        model: String = "glm-4.6",
        temperature: Float = 0.8f,
        maxTokens: Int = 2048
    ): Result<String> = withContext(Dispatchers.IO) {
        try {
            val request = ZhipuChatRequest(
                model = model,
                messages = messages,
                stream = false,
                temperature = temperature,
                maxTokens = maxTokens
            )
            
            val requestBody = gson.toJson(request)
                .toRequestBody("application/json".toMediaType())
            
            val httpRequest = Request.Builder()
                .url("$baseUrl/chat/completions")
                .addHeader("Authorization", "Bearer $apiKey")
                .addHeader("Content-Type", "application/json")
                .post(requestBody)
                .build()
            
            Log.d(TAG, "Sending non-streaming request to Zhipu AI")
            
            val response = client.newCall(httpRequest).execute()
            val responseBody = response.body?.string()
            
            if (!response.isSuccessful || responseBody == null) {
                val errorMsg = responseBody?.let {
                    try {
                        val errorResponse = gson.fromJson(it, ZhipuErrorResponse::class.java)
                        "${errorResponse.error.code}: ${errorResponse.error.message}"
                    } catch (e: Exception) {
                        it
                    }
                } ?: "HTTP ${response.code}"
                
                Log.e(TAG, "Request failed: $errorMsg")
                return@withContext Result.failure(Exception(errorMsg))
            }
            
            val chatResponse = gson.fromJson(responseBody, ZhipuChatResponse::class.java)
            val content = chatResponse.choices.firstOrNull()?.message?.content
                ?: return@withContext Result.failure(Exception("No content in response"))
            
            Log.d(TAG, "Received response: ${content.take(100)}...")
            Result.success(content)
            
        } catch (e: Exception) {
            Log.e(TAG, "Failed to send request", e)
            Result.failure(e)
        }
    }
    
    /**
     * Build conversation history
     */
    fun buildMessages(
        systemPrompt: String? = null,
        conversationHistory: List<Pair<String, String>> = emptyList(),
        userMessage: String
    ): List<ZhipuMessage> {
        val messages = mutableListOf<ZhipuMessage>()
        
        // Add system message if provided
        systemPrompt?.let {
            messages.add(ZhipuMessage(role = "system", content = it))
        }
        
        // Add conversation history
        conversationHistory.forEach { (role, content) ->
            messages.add(ZhipuMessage(role = role, content = content))
        }
        
        // Add current user message
        messages.add(ZhipuMessage(role = "user", content = userMessage))
        
        return messages
    }
}
